from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.edge.service import Service
from selenium.common.exceptions import TimeoutException, WebDriverException
import time
import os
import subprocess
import csv
import re
import json
import winreg
import requests
import zipfile
import shutil
import platform

# ===== 自动下载EdgeDriver的函数 =====
def get_edge_version():
    """从注册表获取Edge浏览器版本"""
    try:
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Edge\BLBeacon")
        version, _ = winreg.QueryValueEx(key, "version")
        winreg.CloseKey(key)
        return version
    except Exception:
        try:
            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\EdgeUpdate")
            version, _ = winreg.QueryValueEx(key, "version")
            winreg.CloseKey(key)
            return version
        except Exception:
            return None

def download_edgedriver(version):
    """下载并解压匹配的EdgeDriver"""
    arch = "win64" if platform.machine().endswith('64') else "win32"
    url = f"https://msedgedriver.azureedge.net/{version}/edgedriver_{arch}.zip"
    driver_dir = os.path.join(os.path.expanduser("~"), "auto_edgedriver")
    os.makedirs(driver_dir, exist_ok=True)
    
    print(f"正在下载EdgeDriver {version} for {arch}...")
    try:
        response = requests.get(url, stream=True, timeout=30)
        if response.status_code == 200:
            zip_path = os.path.join(driver_dir, "edgedriver.zip")
            with open(zip_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(driver_dir)
            os.remove(zip_path)
            
            driver_path = os.path.join(driver_dir, "msedgedriver.exe")
            print(f"EdgeDriver已下载到: {driver_path}")
            return driver_path
    except Exception as e:
        print(f"下载EdgeDriver时出错: {e}")
    return None

def auto_get_edgedriver_path():
    """自动获取匹配的EdgeDriver路径"""
    edge_version = get_edge_version()
    if not edge_version:
        raise Exception("无法获取Edge浏览器版本")
    
    driver_dir = os.path.join(os.path.expanduser("~"), "auto_edgedriver")
    driver_path = os.path.join(driver_dir, "msedgedriver.exe")
    
    if not os.path.exists(driver_path):
        downloaded_path = download_edgedriver(edge_version)
        if downloaded_path:
            return downloaded_path
        else:
            raise Exception("EdgeDriver下载失败")
    return driver_path

# ===== 主程序开始 =====
try:
    edge_driver_path = auto_get_edgedriver_path()
    print(f"使用自动下载的EdgeDriver: {edge_driver_path}")
except Exception as e:
    print(f"自动下载失败: {e}")
    edge_driver_path = r"C:\Users\Devin\.cache\selenium\msedgedriver\win64\136.0.3240.47\msedgedriver.exe"
    print(f"使用备选路径: {edge_driver_path}")

user_data_dir = os.path.expanduser("~") + r"\AppData\Local\Microsoft\Edge\User Data"
profile_dir = "Default"

def clean_browser_resources():
    try:
        subprocess.run(["taskkill", "/F", "/IM", "msedge.exe"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print("已清理浏览器资源")
    except Exception as e:
        print(f"清理浏览器资源时出错: {e}")

clean_browser_resources()

options = webdriver.EdgeOptions()
options.add_argument(f"--user-data-dir={user_data_dir}")
options.add_argument(f"--profile-directory={profile_dir}")

service = Service(executable_path=edge_driver_path)
driver = webdriver.Edge(service=service, options=options)
wait = WebDriverWait(driver, 15)

# 访问抖音个人主页
driver.get("https://www.douyin.com/user/self")
time.sleep(1)

try:
    element = WebDriverWait(driver, 10).until(
        EC.presence_of_element_located((By.CLASS_NAME, 'RlLOO79h'))
    )
    print("登录状态已继承")
except:
    print("登录状态继承失败，请检查用户数据目录")

# 常量定义 - 修改NEW_CSV_PATH使其与CSV_PATH在同一目录
CSV_PATH = r'C:\Users\Devin\Desktop\python\data\comments_cleaned.csv'
csv_dir = os.path.dirname(CSV_PATH)
NEW_CSV_PATH = os.path.join(csv_dir, 'comments_updated.csv')
follow_max = 9999
current_clicks = 0

def convert_count(text):
    if not text or text == '空':
        return 0
    text = text.strip().replace(',', '')
    try:
        if '万' in text:
            return int(float(text.replace('万', '')) * 10000)
        elif '亿' in text:
            return int(float(text.replace('亿', '')) * 100000000)
        return int(text)
    except:
        return 0

# 修正后的header（与清洗后CSV列顺序一致）
header = ["用户名", "位置", "关注数", "粉丝数", "关注按钮状态", "处理情况", "备注", "用户链接"]

# 读取并处理数据
rows = []
with open(CSV_PATH, 'r', encoding='utf-8-sig', errors='replace') as infile:
    reader = csv.reader(infile)
    next(reader)  # 跳过标题行
    for row in reader:
        rows.append(row)
updated_rows = []
total_rows = len(rows)
for i, row in enumerate(rows, 1):
    username = row[0]
    link = row[7]
    progress = (i/total_rows)*100
    print(f"[{progress:.1f}%] 处理用户({i}/{total_rows}): {username}")
    try:
        driver.get(link)
        time.sleep(5)
        # 获取关注按钮状态
        try:
            follow_button = WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, 'span.semi-button-content')))
            button_text = follow_button.text.strip()
            print(f"  关注按钮状态: {button_text}")
            following = WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, 'div.C1cxu0Vq')))
            fans = WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, 'div:nth-child(2) > div.C1cxu0Vq')))
            fans_count = convert_count(fans.text)
            following_count = convert_count(following.text)
            print(f"  关注数: {following_count}, 粉丝数: {fans_count}")
        except TimeoutException:
            print("  🚫 关注数/粉丝数未找到,请退出程序检查页面结构并重新定义选择器")
        # 更新行数据
        row[2] = str(following_count)
        row[3] = str(fans_count)
        row[4] = button_text
        row[5] = "已处理"

    except WebDriverException as e:
        print(f"  🔥 浏览器异常: {str(e)[:50]}")
        row[5] = f"{username}_浏览器异常"
    except Exception as e:
        print(f"  ⛔ 未知错误: {type(e).__name__}")
        row[5] = f"{username}_未知错误"
    
    updated_rows.append(row)
    print("  " + "-"*40)

# 保存结果
with open(NEW_CSV_PATH, 'w', encoding='utf-8-sig', newline='') as outfile:
    writer = csv.writer(outfile)
    writer.writerow(header)
    writer.writerows(updated_rows)

print("\n" + "="*50)
print(f"处理完成:")
print(f"  原始数据: {total_rows} 条")
print(f"  成功处理: {len([r for r in updated_rows if '已处理' in r[5]])} 条")
print(f"  实际关注: {current_clicks} 人")
print(f"  新文件: {NEW_CSV_PATH}")
print("="*50)

driver.quit()